package cn.feizhu.aimarket.ai.workflow.config;

import cn.feizhu.aimarket.ai.client.*;
import cn.feizhu.aimarket.ai.workflow.model.constant.WorkflowControlConstant;
import cn.feizhu.aimarket.ai.workflow.model.enums.NodeStatus;
import cn.feizhu.aimarket.ai.workflow.node.*;
import cn.feizhu.aimarket.ai.workflow.service.WorkflowDataService;
import cn.feizhu.aimarket.tools.MinioUtils;
import com.alibaba.cloud.ai.graph.GraphRepresentation;
import com.alibaba.cloud.ai.graph.KeyStrategy;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.action.AsyncEdgeAction;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.KeyStrategyFactory;
import com.alibaba.cloud.ai.graph.state.strategy.AppendStrategy;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;
import jakarta.annotation.Resource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.HashMap;
import java.util.Map;

import static com.alibaba.cloud.ai.graph.action.AsyncNodeAction.node_async;

/**
 * 智能内容生成工作流配置
 * 流程：智能判断 → 指令处理 → 大纲生成 → 人类反馈 → 多节点生成 → 聚合 → 可视化 → 返回URL
 *
 * @author feizhu
 * @since 2025/10/16
 */
@Configuration
public class GraphConfiguration {

    private static final Logger logger = LoggerFactory.getLogger(GraphConfiguration.class);

    /**
     * 全部使用工厂类创建对话客户端否则会有NPE，因为内部的节点不交给Spring管理
     */
    @Resource
    private HumanResultClientFactory humanResultClientFactory;

    @Resource
    private WriterRouterChatClientFactory writerRouterChatClientFactory;

    @Resource
    private OutLineChatClientFactory outLineChatClientFactory;

    @Resource
    private CommandProcessChatClientFactory commandProcessChatClientFactory;

    @Resource
    private DocGenerateClientFactory docGenerateClientFactory;

    @Resource
    private SimpleAiChatClientFactory simpleAiChatClientFactory;

    @Resource
    private MinioUtils minioUtils;

    @Resource
    private WorkflowDataService workflowDataService;

    /**
     * 对于每一个请求创建单独的contentFlowGraph
     * @param chatClientBuilder
     * @return
     * @throws GraphStateException
     */
    @Bean
    public StateGraph contentFlowGraph(ChatClient.Builder chatClientBuilder) throws GraphStateException {

        // === 节点状态管理 ===
        Map<String, NodeStatus> node2Status = new HashMap<>();

        // === Key 策略 === TODO 这里都要完整注册
        KeyStrategyFactory keyStrategyFactory = () -> {
            HashMap<String, KeyStrategy> map = new HashMap<>();
            map.put(WorkflowControlConstant.USER_INPUT, new ReplaceStrategy());
            map.put(WorkflowControlConstant.JUDGE_RESULT, new ReplaceStrategy()); // TODO 这里之前没有注册导致没有这个
            map.put(WorkflowControlConstant.FEEDBACK, new ReplaceStrategy()); // 人类反馈的结果
            map.put(WorkflowControlConstant.FEEDBACK_RESULT, new ReplaceStrategy()); // AI判断人类反馈的结果
            map.put(WorkflowControlConstant.PROCESSED_COMMAND, new ReplaceStrategy());
            map.put(WorkflowControlConstant.OUTLINE, new ReplaceStrategy());
            map.put(WorkflowControlConstant.SUB_DOCS, new AppendStrategy()); // TODO 这里选择为Append因为子文档要添加一起
            map.put(WorkflowControlConstant.AGGREGATED_CONTENT, new ReplaceStrategy()); // TODO 又忘记注册了
            map.put(WorkflowControlConstant.AGGREGATE_STATS, new ReplaceStrategy()); // TODO 又忘记注册了
            map.put(WorkflowControlConstant.VISUAL_URL, new ReplaceStrategy());
            map.put(WorkflowControlConstant.FINAL_URL, new ReplaceStrategy());
            map.put(WorkflowControlConstant.WORKFLOW_STATUS, new ReplaceStrategy());
            map.put(WorkflowControlConstant.USER_ID, new ReplaceStrategy()); // 用户ID
            map.put(WorkflowControlConstant.APP_ID, new ReplaceStrategy()); // 应用ID
            return map;
        };

        // === 构建工作流 ===
        StateGraph graph = new StateGraph(keyStrategyFactory)
                // 节点注册
                .addNode(WorkflowControlConstant.SMART_CHECK, node_async(new SmartJudgeNode(writerRouterChatClientFactory.createRoutingClient())))
                .addNode(WorkflowControlConstant.COMMAND_HANDLE, node_async(new CommandProcessNode(commandProcessChatClientFactory.createCommandProcessClient())))
                .addNode(WorkflowControlConstant.OUTLINE_GENERATE, node_async(new OutlineGenerateNode(outLineChatClientFactory.createOutLineClient(), node2Status)))
                .addNode(WorkflowControlConstant.OUTLINE_CHECK, node_async(new OutlineCheckNode(simpleAiChatClientFactory.createSimpleAiChatClient(), node2Status)))
                .addNode(WorkflowControlConstant.HUMAN_FEEDBACK, node_async(new HumanFeedbackNode(humanResultClientFactory.createHumanResultClient(), node2Status)))
                .addNode(WorkflowControlConstant.MULTI_GENERATE, node_async(new MultiDocGenerateNode(docGenerateClientFactory, node2Status)))
                .addNode(WorkflowControlConstant.AGGREGATE, node_async(new AggregateNode(node2Status)))
                .addNode(WorkflowControlConstant.VISUALIZE, node_async(new VisualizationNode(minioUtils, workflowDataService, node2Status)))
                .addNode(WorkflowControlConstant.RETURN_URL, node_async(new ReturnUrlNode()))

                // 状态边连接
                .addEdge(StateGraph.START, WorkflowControlConstant.SMART_CHECK)
                .addConditionalEdges(
                        WorkflowControlConstant.SMART_CHECK,
                        AsyncEdgeAction.edge_async(new SmartJudgeDispatcher()),
                        Map.of(
                                WorkflowControlConstant.BRANCH_PASS, WorkflowControlConstant.COMMAND_HANDLE,
                                WorkflowControlConstant.BRANCH_FAIL, StateGraph.END
                        )
                )
                .addEdge(WorkflowControlConstant.COMMAND_HANDLE, WorkflowControlConstant.OUTLINE_GENERATE)
                .addEdge(WorkflowControlConstant.OUTLINE_GENERATE, WorkflowControlConstant.OUTLINE_CHECK)  // 在大纲和人类反馈节点中间新增了一层大纲检查和修复节点
                .addEdge(WorkflowControlConstant.OUTLINE_CHECK, WorkflowControlConstant.HUMAN_FEEDBACK)
                .addConditionalEdges(
                        WorkflowControlConstant.HUMAN_FEEDBACK,
                        AsyncEdgeAction.edge_async(new HumanFeedbackDispatcher()),
                        Map.of(
                                WorkflowControlConstant.BRANCH_CONTINUE, WorkflowControlConstant.MULTI_GENERATE,
                                WorkflowControlConstant.BRANCH_RETRY, WorkflowControlConstant.COMMAND_HANDLE
                        )
                )
                .addEdge(WorkflowControlConstant.MULTI_GENERATE, WorkflowControlConstant.AGGREGATE)
                .addEdge(WorkflowControlConstant.AGGREGATE, WorkflowControlConstant.VISUALIZE)
                .addEdge(WorkflowControlConstant.VISUALIZE, WorkflowControlConstant.RETURN_URL)
                .addEdge(WorkflowControlConstant.RETURN_URL, StateGraph.END);
        //   === 输出 UML 图结构 ===
        GraphRepresentation representation = graph.getGraph(GraphRepresentation.Type.PLANTUML,
                "Content Flow Graph");
        logger.info("\n=== content flow UML ===");
        logger.info(representation.content());
        logger.info("==================================\n");
        return graph;
    }
}
